백준 15681번

오늘 풀어본 문제는 백준의 15681번 문제¹이다. 문제 풀이에 사용한 언어는 C++ 이다.

solved.ac 기준 CLASS

CLASS 5

문제 정보

이 문제의 내용과 조건은 다음과 같다.

문제

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

정점 $U$ 를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

(이 글에서 힌트 부분은 생략하겠다.)

입력

트리의 정점의 수 $N$ 과 루트의 번호 $R$, 쿼리의 수 $Q$ 가 주어진다. $(2 \le N \le 105, 1 \le R \le N, 1 \le Q \le 105)$

이어 $N-1$ 줄에 걸쳐, $U\ V$ 의 형태로 트리에 속한 간선의 정보가 주어진다. $(1 \le U, V \le N, U \neq V)$

이는 $U$ 와 $V$ 를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.

이어 $Q$ 줄에 걸쳐, 문제에 설명한 $U$ 가 하나씩 주어진다. $(1 \le U \le N)$

입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

출력

$Q$ 줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

풀이과정

1번째 시도

문제를 보고 각 노드를 루트 노드로 하는 subtree의 노드 개수를 미리 모두 알아낸 뒤 쿼리에 맞춰 답변하면 되겠다고 생각했다. 이를 위해 재귀 함수를 활용한 DFS를 이용하기로 하였다.

코드는 다음과 같이 작성하였다.

#include <bits/extc++.h>

using namespace std;
using namespace __gnu_pbds;

void setSubTreeNodeCount(int node, int parent, vector<vector<bool>>& graph, vector<int>& dp);

int main(void) {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int N, R, Q;
    cin >> N >> R >> Q;

    vector<vector<bool>> graph(N+1, vector<bool>(N+1, false));

    for (int i=0; i<N-1; i++) {
        int U, V;
        cin >> U >> V;

        graph[U][V] = true;
        graph[V][U] = true;
    }

    vector<int> dp(N+1, 0);

    setSubTreeNodeCount(R, 0, graph, dp);

    for (int i=0; i<Q; i++) {
        int U;
        cin >> U;

        cout << dp[U] << "\n";
    }
     
    return 0;
}

void setSubTreeNodeCount(int node, int parent, vector<vector<bool>>& graph, vector<int>& dp) {
    int N = graph.size() - 1;
    dp[node] = 1;
    
    for (int i=1; i<=N; i++) {
        if (graph[node][i] && i != parent) {
            setSubTreeNodeCount(i, node, graph, dp);
            dp[node] += dp[i];
        }
    }
}

실행 결과 ‘메모리 초과’가 떴다.

2번째 시도

메모리 초과를 해결하기 위해 그래프의 표현 방식을 adjacency matrix 대신 adjaceny list로 바꿔보았다.

코드는 다음과 같이 수정하였다.

#include <bits/extc++.h>

using namespace std;
using namespace __gnu_pbds;

void setSubTreeNodeCount(int node, int parent, vector<vector<int>>& graph, vector<int>& dp);

int main(void) {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int N, R, Q;
    cin >> N >> R >> Q;

    vector<vector<int>> graph(N+1);

    for (int i=0; i<N-1; i++) {
        int U, V;
        cin >> U >> V;

        graph[U].push_back(V);
        graph[V].push_back(U);
    }

    vector<int> dp(N+1, 0);

    setSubTreeNodeCount(R, 0, graph, dp);

    for (int i=0; i<Q; i++) {
        int U;
        cin >> U;

        cout << dp[U] << "\n";
    }

    return 0;
}

void setSubTreeNodeCount(int node, int parent, vector<vector<int>>& graph, vector<int>& dp) {
    int N = graph.size() - 1;
    dp[node] = 1;

    for (int child : graph[node]) {
        if (child != parent) {
            setSubTreeNodeCount(child, node, graph, dp);
            dp[node] += dp[child];
        }
    }
}

그러자 모든 테스트 케이스를 통과하고 ‘맞았습니다’가 나오는 것을 확인할 수 있었다.

마무리

이번 문제는 예전에 B-Tree 를 구현하느라 고생한 경험이 있어서 그런지 생각보다 수월하게 풀이를 구현할 수 있었다. 맨날 이렇게 술술 잘 풀리면 좋겠는데, 앞으로 힘들어지면 더 힘들어졌지 절대 쉬워질 것 같지는 않다…

어쨌든 마지막 CLASS 5 문제를 풀어내면서 솔브드 업데이트로 인해 뺏겼던 CLASS 5 금장을 다시 되찾는데 성공하였다!

CLASS 5 MASTER

이제 진짜 시작이다… 입대전에 CLASS 6, 7 금장을 따는걸 목표로 세웠는데, 상당히 고된 여정이 될 것 같다… 물론 그렇다고 포기할 생각은 없다. 하다보면 감이 생기지 않겠는가?

오늘의 PS는 여기까지!

1: https://www.acmicpc.net/problem/15681

← Previous Post Next Post →